import tensorflow as tf
import numpy as np
from keras.layers import TFSMLayer
from tensorflow.keras.preprocessing.sequence import pad_sequences

# 加载词汇表
word2idx = {}
with open('./vocab/word.txt', encoding='utf-8') as f:
    for i, line in enumerate(f):
        line = line.rstrip()
        word2idx[line] = i

# 定义预处理函数
def preprocess_text(text, max_len=1000):
    # 分词
    words = text.split()
    # 映射到索引
    indices = [word2idx.get(word, len(word2idx)) for word in words]
    # 填充序列
    padded_indices = pad_sequences([indices], maxlen=max_len, padding='post', truncating='post')
    return np.array(padded_indices, dtype=np.float32)  # 改为返回 float32 类型

# 加载 TensorFlow SavedModel 格式的模型
model_layer = TFSMLayer('./models/best_model', call_endpoint='serving_default')

# 将加载的模型层包装成一个 Keras 模型
model = tf.keras.Sequential([model_layer])

# 检查模型结构
model.summary()

# 输入文本
input_text = "I love this movie, it is so amazing!"

# 预处理输入数据
input_data = preprocess_text(input_text)

# 确保输入数据是 TensorFlow 张量，并添加批次维度
input_data = tf.constant(input_data, dtype=tf.float32)  # 使用 float32

# 打印输入数据的类型
print("Input data type:", input_data.dtype)

# 使用模型进行预测
predictions = model.predict(input_data)

# 打印预测结果
print(predictions)
